import pandas as pd
import joblib
import jieba
import re
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import SVC
from sklearn.naive_bayes import MultinomialNB
from sklearn.ensemble import RandomForestClassifier
from sklearn.pipeline import make_pipeline
from sklearn.metrics import classification_report
from backend.utils.db import get_db_connection

# 添加中文停用词表
STOP_WORDS = set(["的", "了", "和", "是", "就", "都", "而", "及", "与", "这", "那", "在", "要"])


class IntentTrainer:
    def preprocess_text(self, text):
        """文本预处理：清洗、分词、移除停用词"""
        # 移除特殊字符和数字
        text = re.sub(r"[^\u4e00-\u9fa5a-zA-Z]", " ", text)
        # 中文分词
        words = jieba.cut(text)
        # 移除停用词和单字词
        words = [word for word in words if word not in STOP_WORDS and len(word) > 1]
        return " ".join(words)

    def balance_dataset(self, df):
        """通过复制少数类样本平衡数据集（不使用SMOTE）"""
        # 计算每个类别的样本数量
        intent_counts = df['intent'].value_counts()
        print("原始类别分布:", intent_counts.to_dict())

        # 找到最大样本数
        max_count = intent_counts.max()

        # 对每个类别进行平衡处理
        balanced_dfs = []
        for intent, group in df.groupby('intent'):
            # 计算需要复制的次数
            current_count = len(group)
            if current_count < max_count:
                # 计算需要复制的倍数和余数
                multiply = max_count // current_count
                remainder = max_count % current_count

                # 复制样本
                repeated = pd.concat([group] * multiply, ignore_index=True)
                # 处理余数
                if remainder > 0:
                    repeated = pd.concat([repeated, group.sample(remainder, random_state=42)], ignore_index=True)
                balanced_dfs.append(repeated)
            else:
                balanced_dfs.append(group)

        # 合并所有类别
        balanced_df = pd.concat(balanced_dfs, ignore_index=True)
        print("平衡后类别分布:", balanced_df['intent'].value_counts().to_dict())
        return balanced_df

    def train_model(self):
        """训练意图分类模型（不使用SMOTE）"""
        # 1. 从数据库加载训练数据
        data = self.load_training_data()
        df = pd.DataFrame(data, columns=["text", "intent"])

        # 2. 文本预处理
        df["cleaned_text"] = df["text"].apply(self.preprocess_text)

        # 3. 平衡数据集（新增步骤）
        df_balanced = self.balance_dataset(df)

        # 4. 划分训练集和测试集（使用平衡后的数据集）
        X_train, X_test, y_train, y_test = train_test_split(
            df_balanced["cleaned_text"], df_balanced["intent"],
            test_size=0.2, random_state=42, stratify=df_balanced["intent"]
        )

        # 5. 特征提取（调整参数）
        vectorizer = TfidfVectorizer(
            max_features=3000,
            ngram_range=(1, 2),  # 包含单字和双字组合
            min_df=3,  # 忽略低频词
            max_df=0.8  # 忽略高频词
        )

        # 6. 尝试不同分类器（添加类别权重参数）
        classifiers = {
            "SVM": SVC(kernel='linear', C=0.5, probability=True, class_weight='balanced'),
            "NaiveBayes": MultinomialNB(class_prior=None),  # 自动调整先验概率
            "RandomForest": RandomForestClassifier(
                n_estimators=100,
                random_state=42,
                class_weight='balanced'  # 自动平衡类别权重
            )
        }

        best_model = None
        best_score = 0
        best_model_name = ""

        print("===== 模型评估结果 =====")
        for name, clf in classifiers.items():
            # 创建管道
            pipeline = make_pipeline(vectorizer, clf)

            # 交叉验证
            cv_scores = cross_val_score(pipeline, X_train, y_train, cv=5)
            cv_mean = cv_scores.mean()

            # 完整训练
            pipeline.fit(X_train, y_train)
            test_acc = pipeline.score(X_test, y_test)

            print(f"{name}:")
            print(f"  交叉验证准确率: {cv_mean:.4f}")
            print(f"  测试集准确率: {test_acc:.4f}")

            # 保存最佳模型
            if test_acc > best_score:
                best_model = pipeline
                best_score = test_acc
                best_model_name = name

        # 7. 输出最佳模型报告
        y_pred = best_model.predict(X_test)
        print("\n===== 最佳模型分类报告 =====")
        print(classification_report(y_test, y_pred))

        # 8. 保存最佳模型
        joblib.dump(best_model, "intent_model.pkl")
        print(f"保存最佳模型: {best_model_name}")

        # 保存向量器
        joblib.dump(vectorizer, "tfidf_vectorizer.pkl")

        return best_model

    def load_training_data(self):
        """从数据库加载训练数据（增加默认数据量）"""
        data = []
        try:
            with get_db_connection() as (cursor, _):
                cursor.execute("SELECT text, intent FROM intent_training_data")
                for row in cursor.fetchall():
                    data.append((row["text"], row["intent"]))
        except Exception as e:
            print(f"加载训练数据失败: {e}")

        # 增强默认数据集（覆盖更多场景）
        if not data or len(data) < 100:  # 增加到至少100条
            default_data = [
                # product_query (产品查询) - 增加更多变体
                ("这个多少钱", "product_query"),
                ("价格是多少", "product_query"),
                ("卖多少钱", "product_query"),
                ("怎么卖", "product_query"),
                ("价格多少", "product_query"),
                ("售价", "product_query"),
                ("报价", "product_query"),
                ("什么价格", "product_query"),
                ("多少钱一个", "product_query"),
                ("价格贵吗", "product_query"),
                ("这款价格", "product_query"),
                ("价格优惠吗", "product_query"),
                ("最低价多少", "product_query"),
                ("现在什么价", "product_query"),
                ("有折扣吗", "product_query"),
                ("促销价多少", "product_query"),
                ("会员价多少", "product_query"),
                ("最终价格", "product_query"),
                ("多少钱能买", "product_query"),
                ("价格范围", "product_query"),

                # order_query (订单查询) - 增加更多变体
                ("查订单", "order_query"),
                ("订单状态", "order_query"),
                ("发货了吗", "order_query"),
                ("到哪了", "order_query"),
                ("物流信息", "order_query"),
                ("快递到哪", "order_query"),
                ("订单进度", "order_query"),
                ("查看我的订单", "order_query"),
                ("订单号查询", "order_query"),
                ("查一下订单", "order_query"),
                ("订单情况", "order_query"),
                ("订单跟踪", "order_query"),
                ("我的订单", "order_query"),
                ("订单发货没", "order_query"),
                ("订单到哪里了", "order_query"),
                ("查物流", "order_query"),
                ("快递信息", "order_query"),
                ("配送状态", "order_query"),
                ("订单处理进度", "order_query"),
                ("订单状态查询", "order_query"),

                # product_skin_type 意图样本
                ("P001适合适合什么肤质", "product_skin_type"),
                ("这款这款产品适合适用干皮吗", "product_skin_type"),
                ("敏感肌能用这个商品吗", "product_skin_type"),
                ("油皮适合哪款产品", "product_skin_type"),
                ("这个化妆品的适用肤质是什么", "product_skin_type"),

                # product_series 意图样本
                ("P002 属于哪个系列", "product_series"),
                ("这款商品是什么系列的", "product_series"),
                ("这个产品属于特安修护系列吗", "product_series"),
                ("查询产品的系列名称", "product_series"),
                ("这个化妆品属于哪个系列", "product_series"),
                ("P003 属于哪个产品系列", "product_series"),
                ("这款护肤品是多元优效系列吗", "product_series"),
                ("查询P005的所属系列", "product_series"),

                # after_sales_query 意图样本
                ("ORD567 过期了能退吗", "after_sales_query"),
                ("这个商品质量有问题，要退款", "after_sales_query"),
                ("收到的商品变质了，申请退货", "after_sales_query"),
                ("ORD123 产品变质了要退货", "after_sales_query"),
                ("订单过期了能退款吗", "after_sales_query"),
                ("质量问题想退货", "after_sales_query"),
                ("申请退货退款", "after_sales_query"),
                ("这个订单的商品坏了要退", "after_sales_query"),
                ("ORD456 商品破损要换货", "after_sales_query"),
                ("产品有瑕疵能更换吗", "after_sales_query"),
                ("想换一个新的，订单号是ORD789", "after_sales_query"),
                ("申请换货处理", "after_sales_query"),
                ("这个商品有问题，能换同款吗", "after_sales_query"),
                ("退货", "after_sales_query"),
                ("我要退款", "after_sales_query"),
                ("质量问题", "after_sales_query"),
                ("退货流程", "after_sales_query"),
                ("申请退货", "after_sales_query"),
                ("怎么退货", "after_sales_query"),
                ("退货怎么处理", "after_sales_query"),
                ("退换货", "after_sales_query"),
                ("退货政策", "after_sales_query"),
                ("如何退货", "after_sales_query"),
                ("退货申请", "after_sales_query"),
                ("想退货", "after_sales_query"),
                ("退货步骤", "after_sales_query"),
                ("退款流程", "after_sales_query"),
                ("退货退款", "after_sales_query"),
                ("退商品", "after_sales_query"),
                ("退货怎么办", "after_sales_query"),
                ("退货事宜", "after_sales_query"),
                ("退货服务", "after_sales_query"),
                ("退货处理", "after_sales_query"),

                # address_query (修改地址地址) - 增加更多变体
                ("改地址", "address_query"),
                ("修改收货地址", "address_query"),
                ("地址错了", "address_query"),
                ("变更地址", "address_query"),
                ("地址变更", "address_query"),
                ("改收货信息", "address_query"),
                ("换地址", "address_query"),
                ("地址怎么改", "address_query"),
                ("改一下地址", "address_query"),
                ("地址修改", "address_query"),
                ("改收货地址", "address_query"),
                ("修改配送地址", "address_query"),
                ("地址不对", "address_query"),
                ("改送货地址", "address_query"),
                ("更新地址", "address_query"),
                ("地址需要改", "address_query"),
                ("改下地址", "address_query"),
                ("地址能改吗", "address_query"),
                ("改地址信息", "address_query"),
                ("改一下收货地址", "address_query"),

                # human_service (人工服务) - 增加更多变体
                ("转人工", "human_service"),
                ("人工客服", "human_service"),
                ("找真人", "human_service"),
                ("人工服务", "human_service"),
                ("联系客服", "human_service"),
                ("在线客服", "human_service"),
                ("转接人工", "human_service"),
                ("人工帮助", "human_service"),
                ("找人工", "human_service"),
                ("人工支持", "human_service"),
                ("客服专员", "human_service"),
                ("人工处理", "human_service"),
                ("人工坐席", "human_service"),
                ("人工咨询", "human_service"),
                ("人工对话", "human_service"),
                ("人工接入", "human_service"),
                ("人工服务台", "human_service"),
                ("转人工客服", "human_service"),
                ("找客服人员", "human_service"),
                ("人工助手", "human_service"),

                # general_query (通用查询) - 增加更多变体
                ("你好", "general_query"),
                ("谢谢", "general_query"),
                ("再见", "general_query"),
                ("您好", "general_query"),
                ("在吗", "general_query"),
                ("有人吗", "general_query"),
                ("功能说明", "general_query"),
                ("帮助", "general_query"),
                ("嗨", "general_query"),
                ("哈喽", "general_query"),
                ("早上好", "general_query"),
                ("下午好", "general_query"),
                ("晚上好", "general_query"),
                ("拜拜", "general_query"),
                ("感谢", "general_query"),
                ("谢了", "general_query"),
                ("不客气", "general_query"),
                ("请帮忙", "general_query"),
                ("求助", "general_query"),
                ("请问", "general_query"),
                # 补充
                ("你们家有什么乳", "product_query"),
                ("推荐几款好用的乳", "product_query"),
                ("保湿乳多少钱", "product_query"),
                ("修护乳适合什么肤质", "product_skin_type"),
                ("这款乳的成分是什么", "product_ingredients"),
                ("推荐一下乳液", "product_query"),
                ("你们家有什么面霜", "product_query"),
                ("乳液有哪些款式", "product_query"),
                ("面霜的价格是多少", "product_query"),
                ("保湿乳液适合什么肤质", "product_skin_type"),
                ("这款面霜的成分是什么", "product_ingredients"),
                ("推荐一款好用的面霜", "product_query"),
                ("乳液和面霜有什么区别", "product_query"),
                ("你们的明星产品是哪款乳液", "product_query"),
                ("面霜有促销活动吗", "product_query"),
                ("我的乳液有质量问题要退货", "after_sales_return"),
                ("面霜过敏了能换货吗", "after_sales_exchange"),
                ("订单里的乳液变质了", "after_sales_return"),
                ("P001多少钱", "product_query"),
                ("净肌平衡清透水价格是多少", "product_query"),
                ("特安修护乳卖多少钱", "product_query"),
                ("ORD123到哪了", "order_query"),
                ("订单ORD456发货了吗", "order_query"),
                ("ORD123想退货", "after_sales_query"),
                ("订单产品变质能退款吗", "after_sales_query"),
                ("P001的成分是什么", "product_ingredients"),
                ("净肌平衡清透水有哪些配料", "product_ingredients"),
                ("P002有什么功效", "product_effect"),
                ("特安修护乳效果怎么样", "product_effect"),
                ("转人工客服", "human_service"),
                ("这个问题找客服解决", "human_service")
            ]
            data.extend(default_data)

        return data


if __name__ == "__main__":
    # 初始化jieba分词
    jieba.initialize()

    # # 设置matplotlib中文字体
    # plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

    trainer = IntentTrainer()
    trainer.train_model()